package edu.purdue.cs505;

import java.io.IOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetAddress;
import java.net.SocketTimeoutException;
import java.net.UnknownHostException;
import java.util.Collections;
import java.util.Date;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.PriorityBlockingQueue;

/**
 * Implementation of the ReliableChannel interface.
 */
public class ReliableChannelImpl implements ReliableChannel
{
  /*
   * How this works: the priority queue contains all outstanding messages ordered by timeout. When
   * checking timeout, we check the head of the queue. If it is expired, we then check the ackedIds
   * set to see if it has been acked. If it has, we remove from the queue and check the new head. If
   * it hasn't, we resend it and change its timeout (and it's position in the queue). We continue
   * checking the new head until the queue is empty or we hit an unexpired timeout.
   */
  private Set<Long> ackedIds;
  private PriorityBlockingQueue<DataMessage> outstandingMessages;
  private long sendId;
  private DatagramSocket socket;
  private InetAddress localSocketIP;
  private int localSocketPort;
  private InetAddress destSocketIP;
  private int destSocketPort;
  private String myProcessId;

  // receiver things:
  private Set<Long> receivedOutOfOrder;
  private long receivedAllBelow;
  private ReliableChannelReceiver rc;

  boolean halted;

  public ReliableChannelImpl(String socketIP, int socketPort)
  {
    try {
      this.localSocketIP = InetAddress.getByName(socketIP);
    }
    catch (UnknownHostException e) {
      e.printStackTrace();
    }
    this.localSocketPort = socketPort;
  }

  /**
   * Start up the socket, set counters, initialize data structures
   */
  public void init(String destinationIP, int destinationPort)
  {
    ackedIds = Collections.synchronizedSet(new HashSet<Long>());
    outstandingMessages = new PriorityBlockingQueue<DataMessage>();
    receivedOutOfOrder = Collections.synchronizedSet(new TreeSet<Long>()); // TreeSet
    // for ordering of elements
    rc = null;
    receivedAllBelow = 0;
    sendId = 0;
    halted = false;

    // set up the sockets
    try {
      socket = new DatagramSocket(localSocketPort, localSocketIP);
      socket.setSoTimeout(Constants.SOCKET_TIMEOUT);
      destSocketIP = InetAddress.getByName(destinationIP);
      destSocketPort = destinationPort;
      //System.out.println("local: " + getLocalId() + " remote: " + getDestinationId());
    }
    catch (IOException e) {
      e.printStackTrace();
    }

    MessageTimeoutChecker runnable = new MessageTimeoutChecker();
    new Thread(runnable).start();
    SocketListener listener = new SocketListener();
    new Thread(listener).start();
  }

  public void setMyProcessId(String processId)
  {
    this.myProcessId = processId;
  }

  /**
   * Sends the initial message, adding it to the outstanding list.
   */
  public void rsend(Message m)
  {
    /* only update the send id if we originate a new message */
    DataMessage dm;
    if (m.getProcessID().equals(this.myProcessId)) {
      sendId++;
      dm = new DataMessage(m, sendId);
    } else {
      dm = new DataMessage(m, m.getMessageNumber());
    }
    dm.setTimeout(Constants.MESSAGE_TIMEOUT);
    synchronized (this) {
    outstandingMessages.add(dm);
    }
    send(dm);
  }

  /**
   * Simulate a dropped packet by failing to send the message over the network.
   */
  public void rsendFail(Message m)
  {
    DataMessage dm = new DataMessage(m, sendId++);
    dm.setTimeout(Constants.MESSAGE_TIMEOUT);
    synchronized (this) {
    outstandingMessages.add(dm);
    }
  }

  /**
   * Actually sends the message on the socket, used for both the initial send and any retries.
   */
  private void send(DataMessage dm)
  {
    // //System.out.println("Sending data: " + dm);
    dm.setTimeout(Constants.MESSAGE_TIMEOUT);
    // serialize & send it on the socket
    byte[] serialized = dm.serialize().getBytes();
    try {
      socket.send(new DatagramPacket(serialized, serialized.length, destSocketIP, destSocketPort));
    }
    catch (IOException e) {
      e.printStackTrace();
    }
  }

  /**
   * Send ACK to socket
   */
  private void send(AckMessage ack)
  {
    //System.out.println("Sending ack: " + ack);
    // serialize & send it on the socket
    byte[] serialized = ack.serialize().getBytes();
    try {
      socket.send(new DatagramPacket(serialized, serialized.length, destSocketIP, destSocketPort));
    }
    catch (IOException e) {
      e.printStackTrace();
    }
  }

  /**
   * Specify a listener for this channel
   */
  public void rlisten(ReliableChannelReceiver rc)
  {
    this.rc = rc;
  }

  /**
   * Stop this channel
   */
  public void halt()
  {
    halted = true;
  }

  public DatagramSocket getSocket()
  {
    return socket;
  }

  private String getDestinationId()
  {
    return destSocketIP.getHostAddress() + ":" + destSocketPort;
  }
  
  private String getLocalId()
  {
    return localSocketIP.getHostAddress() + ":" + localSocketPort;
  }
  
  /* thread that checks timeouts on send messages */
  private class MessageTimeoutChecker implements Runnable
  {
    public void run()
    {
      while (!halted) {
        /*
         * look at the message with the closest timeout. if it is timed out, remove it from the
         * priority queue and send it again. then add it back to the priority queue
         */
        synchronized (this) {
        if (outstandingMessages.size() > 0) {
          DataMessage dm = outstandingMessages.peek();
          long now = new Date().getTime();
          while (dm != null && now >= dm.getTimeout()) {
            // remove that message
            outstandingMessages.poll();
            if (!ackedIds.contains(dm.getId())) {
              // send it again
              send(dm);
              //System.out.println("Resending (" + getDestinationId() + "): " + dm);
              // add it back to the priority queue with a new
              // timeout
              dm.setTimeout(Constants.MESSAGE_TIMEOUT);
              outstandingMessages.add(dm);
            }
            else {
              // clean this up since it won't be relevant again
              ackedIds.remove(dm.getId());
            }
            dm = outstandingMessages.peek();
          }
        }
        }
      }
      // //System.out.println("Timeout checker halted");
    }
  }

  /* thread that listens to the UDP socket for new datagrams */
  private class SocketListener implements Runnable
  {
    public void run()
    {
      while (!halted) {
        synchronized (this) {
        DatagramPacket p =
            new DatagramPacket(new byte[Constants.RECEIVE_SIZE], Constants.RECEIVE_SIZE);
        try {
          socket.receive(p);

          String received = new String(p.getData());

          // figure out if this is a DataMessage or an AckMessage and
          // handle properly
          if (AckMessage.isAck(received)) {
            AckMessage ack = AckMessage.deserialize(received);
            //System.out.println("Received ack(" + getDestinationId() + "): " + ack);
            ackedIds.add(ack.getId());
          }
          else if (rc != null) {
            DataMessage dm = DataMessage.deserialize(received);
             //System.out.println("Received data(" + getDestinationId() + "): " + dm);
            long message_id = dm.getId();
            if (message_id == receivedAllBelow) {
              // update which id we expect to receive next
              receivedAllBelow = message_id + 1;
              // receive message
              Message m = dm.getMessage();
              rc.rreceive(m);
              // send ACK

              //System.out.println("Sending ack(" + getDestinationId() + ")(in order): " + message_id);
              send(new AckMessage(message_id));
              // check if we can update expected id from messages
              // we have received out of
              // order
              // the id's should already be in sorted order
              Iterator<Long> iter = receivedOutOfOrder.iterator();
              while (iter.hasNext()) {
                Long id = iter.next();
                assert id >= receivedAllBelow; // we only add
                // id's greater
                // than the
                // current
                // expected id
                if (id == receivedAllBelow) {
                  receivedAllBelow = id + 1;
                  iter.remove();
                }
                else {
                  break; // we cannot update the expected id
                         // anymore
                }
              }
            }
            else if (message_id < receivedAllBelow) {
              // discard message and send ack
              //System.out.println("Sending ack(" + getDestinationId() + ")(discard): " + message_id);
              send(new AckMessage(message_id));
            }
            else if (message_id > receivedAllBelow) {
              // keep track of id's we have received out of order
              receivedOutOfOrder.add(message_id);
              // receive message
              Message m = dm.getMessage();
              rc.rreceive(m);
              // send ACK
              //System.out.println("Sending ack(" + getDestinationId() + ")(OOO): " + message_id);
              send(new AckMessage(message_id));
            }
          }
        }
        catch (SocketTimeoutException e) {
          // receive() timed out, checked halted again and continue
        }
        catch (IOException e) {
          e.printStackTrace();
        }
      }
      // //System.out.println("Socket listener halted");
      }
    }
  }
}
